import torch
import torchvision
from torch import nn
from torch.nn import functional as f
from d2l import torch as d2l

pretrained_net = torchvision.models.resnet18(pretrained=True)
net = nn.Sequential(*list(pretrained_net.children())[:-2])
X = torch.rand(size=(1,3,320,480))
print(net(X).shape)

num_classes = 21
net.add_module('final_conv',nn.Conv2d(512,num_classes,kernel_size=1))
net.add_module('transposed_conv',nn.ConvTranspose2d(num_classes,num_classes,
                                                    kernel_size=64,padding=16,stride=32))
#双线性插值
def bilinear_kernel(in_channels,out_channels,kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = (torch.arange(kernel_size).reshape(-1,1),
          torch.arange(kernel_size).reshape(1.-1))
    filt = (1 - torch.abs(og[0] - center) / factor,
            1 - torch.abs(og[1] - center) / factor)
    weight = torch.zeros((in_channels,out_channels,kernel_size,kernel_size))
    weight[range(in_channels),range(out_channels),:,:] = filt
    return weight

batch_size = 32
crop_size = (320,480)
train_iter,test_iter = d2l.load_data_voc(batch_size,crop_size)
def loss(inputs,targets):
    return f.cross_entropy(inputs,targets,reduction='none').mean(1).mean(1)
num_epochs = 5
lr = 0.001
wd = 1e-3
devices = d2l.try_all_gpus()
trainer = torch.optim.SGD(net.parameters(),lr=lr,weight_decay=wd)
d2l.train_ch13(net,train_iter,test_iter,loss,trainer,num_epochs)

def predict(img):
    X = test_iter.dataset.normalize_images(img).unsqueeze(0)
    pred = net(X.to(devices[0])).argmax(dim=1)
    return pred.reshape(pred.shape[1],pred.shape[2])